from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

class RerankService:
    model = None
    tokenizer = None
    
    @classmethod
    def init_model(cls):
        if cls.model is None:
            cls.model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
            cls.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-base')
    
    @classmethod
    def rerank(cls, query, documents, top_k=None):
        cls.init_model()
        
        pairs = [[query, doc.content] for doc in documents]
        features = cls.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            return_tensors='pt',
            max_length=512
        )
        
        with torch.no_grad():
            scores = cls.model(**features).logits.squeeze()
            
        ranked_indices = torch.argsort(scores, descending=True)
        
        if top_k:
            ranked_indices = ranked_indices[:top_k]
            
        return [documents[idx] for idx in ranked_indices] 